package com.stress.testing.plugin.apache.rocketmq.interceptor;

import com.stress.testing.common.StressTesting;
import com.stress.testing.common.interceptor.AbstractInterceptor;
import com.stress.testing.plugin.apache.rocketmq.context.PushConsumerContext;
import com.stress.testing.plugin.apache.rocketmq.context.PushConsumerContextHolder;
import com.stress.testing.plugin.apache.rocketmq.listener.MessageListenerConcurrentlyWrapper;
import com.stress.testing.plugin.apache.rocketmq.listener.MessageListenerOrderlyWrapper;
import com.stress.testing.plugin.apache.rocketmq.listener.MessageListenerWrapper;
import org.apache.rocketmq.client.consumer.DefaultMQPushConsumer;
import org.apache.rocketmq.client.consumer.listener.MessageListenerConcurrently;
import org.apache.rocketmq.client.consumer.listener.MessageListenerOrderly;

import static com.stress.testing.common.conf.PluginConfigKey.DEFAULT_PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX;
import static com.stress.testing.common.conf.PluginConfigKey.PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX_KEY;

public class DefaultMQPushConsumerRegisterInterceptor extends AbstractInterceptor {

    @Override
    public Object after(Object result, Object[] args, Object zuper, Object thiz) {
        if (!StressTesting.enable(getAgentConfig(), getPluginConfig())) {
            return result;
        }

        if (args == null || args.length == 0) {
            return result;
        }

        if (thiz == null || !(thiz instanceof DefaultMQPushConsumer)) {
            return result;
        }

        // 在改变原始consumer的MessageListener时，返回
        Object arg = args[0];
        if (arg instanceof MessageListenerWrapper) {
            return result;
        }

        // 压测的consumer在new的时候，已经将ConsumerGroup设置了
        DefaultMQPushConsumer consumer = (DefaultMQPushConsumer) thiz;
        String consumerGroup = consumer.getConsumerGroup();
        String topicPrefix = getPluginConfig().get(PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX_KEY,
                DEFAULT_PLUGIN_APACHE_ROCKETMQ_TOPIC_PREFIX);
        if (consumerGroup != null && consumerGroup.startsWith(topicPrefix)) {
            return result;
        }

        PushConsumerContext context = PushConsumerContextHolder.getPushConsumerContext(consumer.hashCode());
        if (context == null) {
            context = new PushConsumerContext();
            PushConsumerContextHolder.setPushConsumerContext(consumer.hashCode(), context);
        }

        context.setRegisterArgs(args);

        // 将原始MessageListener改掉
        if (arg instanceof MessageListenerConcurrently) {
            MessageListenerConcurrently target = (MessageListenerConcurrently) arg;
            MessageListenerConcurrentlyWrapper listener =
                    new MessageListenerConcurrentlyWrapper(getAgentConfig(), getPluginConfig(), target);
            consumer.registerMessageListener(listener);
        } else if (arg instanceof MessageListenerOrderly) {
            MessageListenerOrderly target = (MessageListenerOrderly) arg;
            MessageListenerOrderlyWrapper listener =
                    new MessageListenerOrderlyWrapper(getAgentConfig(), getPluginConfig(), target);
            consumer.registerMessageListener(listener);
        }

        return result;
    }
}
